package com.buaa.pm.service.impl;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.buaa.pm.config.activemq.ActivemqQueueConstant;
import com.buaa.pm.entity.*;
import com.buaa.pm.enumeration.ExtractionStatus;
import com.buaa.pm.feign.ExtractClient;
import com.buaa.pm.mapper.TaskMapper;
import com.buaa.pm.model.dto.ParamRequestDTO;
import com.buaa.pm.model.dto.TaskInfoDTO;
import com.buaa.pm.service.*;
import org.checkerframework.checker.units.qual.A;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jms.annotation.JmsListener;
import org.springframework.stereotype.Service;
import org.springframework.jms.core.JmsMessagingTemplate;
import org.springframework.util.StringUtils;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Service
public class TaskServiceImpl extends ServiceImpl<TaskMapper, Task> implements TaskService {

    @Autowired
    private ParagraphService paragraphService;

    @Autowired
    private ExtractResultService extractResultService;

    @Autowired
    private JmsMessagingTemplate jmsMessagingTemplate;

    @Autowired
    private ExtractClient extractClient;

    @Autowired
    private ExtractRecordService extractRecordService;

    @Autowired
    private PaperService paperService;

    @Override
    public Task createTask(TaskInfoDTO taskInfoDTO) {
        Task task = new Task();
        BeanUtils.copyProperties(taskInfoDTO, task);
        if (!StringUtils.hasText(task.getName())){
            DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyyMMddHHmmssSSS");
            task.setName(LocalDateTime.now().format(dtf));
        }

        // 获取选中的合成段
        QueryWrapper<Paragraph> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(Paragraph::getProjectId, task.getProjectId());
        queryWrapper.lambda().eq(Paragraph::getAvailable,1);
        List<Paragraph> paragraphList = paragraphService.list(queryWrapper);
        task.setParaCount(paragraphList.size());

        // 计算论文数量
        List<String> paperIds = paragraphList.stream()
                .map(Paragraph::getPaperId)
                .distinct()  // 如果一个paperID在paragraphList中出现了多次，只计算一次
                .collect(Collectors.toList());
        QueryWrapper<Paper> queryWrapperPaper = new QueryWrapper<>();
        queryWrapperPaper.lambda()
                .in(Paper::getId, paperIds);
        task.setPaperCount((int) paperService.count(queryWrapperPaper));
        save(task);

        for (Paragraph paragraph : paragraphList) {
            ExtractResult extractResult = new ExtractResult();
            extractResult.setTaskId(task.getId());
            extractResult.setParagraphId(paragraph.getId());
            extractResult.setStatus(ExtractionStatus.WAITING);
            if(extractResultService.save(extractResult)){
                jmsMessagingTemplate.convertAndSend(ActivemqQueueConstant.EXTRACT_QUEUE, extractResult);
            }
        }
        return task;
    }

    @JmsListener(destination = ActivemqQueueConstant.EXTRACT_QUEUE)
    public void receiveExtractResult(ExtractResult extractResult) {
        Task task = getById(extractResult.getTaskId());
        Paragraph paragraph = paragraphService.getById(extractResult.getParagraphId());

        // 设置请求参数
        ParamRequestDTO paramRequestDTO = new ParamRequestDTO();
        paramRequestDTO.setSize(task.getSize());
        paramRequestDTO.setModel(task.getModel());
        paramRequestDTO.setTemperature(task.getTemperature());
        paramRequestDTO.setParagraph(paragraph.getContent());
        paramRequestDTO.setSampling(task.getSampling());

        //发送合成参数抽取请求
        try{
            JSONObject jsonObject = extractClient.getParam(paramRequestDTO);
            // 返回为msg错误消息时会在遍历时抛出异常进入catch
            // 取result
            JSONArray resultArray = jsonObject.getJSONArray("result");
            extractResult.setResult(jsonObject.toJSONString());
            // 取prompt
            extractResult.setPrompt(jsonObject.getJSONArray("prompt").toJSONString());
            int count = 0; // 初始化计数器
            int tagsCount = 0; // 初始化标签计数器
            // 生成ExtractRecord
            //第二次标注数据格式(假设只有一组合成参数)
            for(Object chemicalGroup : resultArray){
                String groupString = JSONObject.toJSONString(chemicalGroup);
                JSONObject group = JSONObject.parseObject(groupString);
                // 遍历标签
                for (String tag : group.keySet()){
                    // 有嵌套的标签
                    if (tag.equals("Metal_Source") || tag.equals("Solvent") || tag.equals("Organic_Linker") || tag.equals("Modulator")) {
                        JSONArray jsonArray = group.getJSONArray(tag);
                        if(jsonArray.size()>0) {
                            tagsCount++;
                        }
                        for (int i = 0; i < jsonArray.size(); i++) {
                            JSONObject innerObj = jsonArray.getJSONObject(i);
                            String precursorName = innerObj.getString("precursor_name");
                            String amount = innerObj.getString("amount");
                            // amount
                            ExtractRecord amountRecord = new ExtractRecord();
                            amountRecord.setContent(amount);
                            amountRecord.setTag(tag);
                            amountRecord.setTaskId(task.getId());
                            amountRecord.setParagraphId(paragraph.getId());
                            extractRecordService.save(amountRecord);
                            count++;
                            String amountId = amountRecord.getId();
                            // precursor_name
                            ExtractRecord precursorNameRecord = new ExtractRecord();
                            precursorNameRecord.setContent(precursorName);
                            precursorNameRecord.setTag(tag);
                            precursorNameRecord.setTaskId(task.getId());
                            precursorNameRecord.setParagraphId(paragraph.getId());
                            precursorNameRecord.setAmount(amountId);
                            extractRecordService.save(precursorNameRecord);
                            count++;
                        }
                    }
                    else{
                        JSONArray contentList = group.getJSONArray(tag);
                        if(contentList.size()>0) {
                            tagsCount++;
                        }
                        for (int i = 0; i < contentList.size(); i++){
                            String content = contentList.getString(i);
                            ExtractRecord extractRecord = new ExtractRecord();
                            extractRecord.setContent(content);
                            extractRecord.setTag(tag);
                            extractRecord.setTaskId(task.getId());
                            extractRecord.setParagraphId(paragraph.getId());
                            extractRecordService.save(extractRecord);
                            count++;
                        }
                    }

                }
            }
            extractResult.setParaLength(paragraph.getContent().length());
            extractResult.setTags(tagsCount);
            extractResult.setCount(count);
            //更新状态
            extractResult.setStatus(ExtractionStatus.FINISH);
            extractResultService.updateById(extractResult);
        }
        catch (Exception e){
            e.printStackTrace();
            //更新状态
            extractResult.setStatus(ExtractionStatus.FAILED);
            extractResultService.updateById(extractResult);
        }
    }

    @Override
    public List<Task> listTaskByProject(String projectId) {
        QueryWrapper<Task> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(Task::getProjectId, projectId);
        queryWrapper.lambda().orderByDesc(Task::getCreatedAt);
        return list(queryWrapper);
    }

    @Override
    public Task getTaskById(String taskId){
        QueryWrapper<Task> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(Task::getId, taskId);
        return list(queryWrapper).get(0);
    }

}
